#include <iostream>
#include <string>
#include <vector>
#include <filesystem>
#include <iomanip>
#include <cmath>
#include <thread>
#include <fstream>
#include <omp.h>
#include "DataTypes.h"
#include "FileIO.h"
#include "Optimizer.h"
#include "CoreLogic.h"

int main() {
    unsigned int max_threads = std::thread::hardware_concurrency();
    unsigned int num_threads = (max_threads > 1) ? max_threads - 1 : 1;
    omp_set_num_threads(num_threads);
    std::cout << "--- Parallel computing enabled with " << num_threads << " threads (Max hardware threads: " << max_threads << ") ---" << std::endl;

    // ========================================================================
    // --- ���������� (Centralized Parameter Settings) ---
    // ========================================================================
    const std::string case_name = "FlashAttention_Case0";
    std::cout << "--- Processing case: " << case_name << " ---" << std::endl;

    // 1. Ӳ���������� (Caps)
    // ----------------------------------------------------
    // ����ģ���и����������������λΪ�ֽڡ�
    Caps caps;
    caps.l1 = 4096;   // L1 �����������
    caps.ub = 1024;   // UB �����������
    caps.l0a = 256;   // L0A �����������
    caps.l0b = 256;   // L0B �����������
    caps.l0c = 512;   // L0C �����������

    // 2. ��ʼ���Կ��Ʋ��� (Params)
    // ----------------------------------------------------
    // ��Щ��������Spill�����㷨����Ϊ���Ż������ڽ��������ж���Щ��������΢����
    Params params;
    // [ʹ��˵��] ����Ҫ�ڳ��ռ�ʱ���˲�������ɸѡ�����롱������Ʒ(victim)��
    // ֻ����δ��������������Żᱻ�ٴ�ʹ�õ�Buffer���Żᱻ���������ѡ��
    // - �����ֵ����������ڱ���������ʹ�õ�Buffer�����ܻᵼ��SpillһЩ�������ںܳ���Buffer������Spill������
    // - ��С��ֵ���������Spill������ʹ�õ�Buffer�����ܽ������ܡ�
    params.min_farthest_use_threshold = 500;

    // [ʹ��˵��] ��������Ȩ��(w_*)�����ڶ�������롱Spill�����������ߣ�����֮�ͽ���Ϊ1.0��
    // w_use: ��δ��ʹ�þ��롱��Ȩ�ء�ֵԽ�ߣ�Խ������ѡ��δ�������Żᱻ�õ���Buffer����Spill (�����Belady�㷨˼��)��
    params.w_use = 0.5;
    // w_vol: ��Spill��������Ȩ�ء�ֵԽ�ߣ�Խ������ѡ���ܴ�С(volume)��С��Buffer��Ͻ���Spill���Լ������ݰ�������
    params.w_vol = 0.3;
    // w_vic: ������Ʒ��������Ȩ�ء�ֵԽ�ߣ�Խ������ѡ������Buffer�������ٵķ������Լ���Spill�¼��ķ���������
    params.w_vic = 0.2;

    // �ڲ������������������ֲ�ͬ�ķ�����ԡ��Ż����ᶯ̬�ı������˴����õ�ֵ�������㷨��ʼʱ�ĵ�һ��������
    params.alloc_strategy = "BestFit";

    // 3. �Ż���ѡ�� (Options)
    // ----------------------------------------------------
    // ����NSGA-II�Ŵ��㷨��ȫ����Ϊ��
    Options options;
    // [ʹ��˵��] �Ż����ܴ�����������������
    // - �����ֵ�������Ż�������ʱ��Ѱ�Ҹ��õĽ⣬���ܺ�ʱ���ӡ�
    // - ��С��ֵ�����ٵõ�����������ܽ���������ߡ�
    options.generations = 50;

    // [ʹ��˵��] ÿһ���б�������Ⱥ��С��
    // - �����ֵ��������Ⱥ�����ԣ������ҵ����õ�������ǰ�أ���ÿ���ļ��������������ӡ�
    options.population_size =10;

    // [ʹ��˵��] ÿһ���Ӹ��������ɵ����Ӵ����½⣩��������
    options.batch_size =10;

    // [ʹ��˵��] ��������ĸ���ѡ��
    // ֵΪ0.8��ζ����80%�ĸ��ʶԡ�����˳�򡱽��б��죬��20%�ĸ��ʶԡ�Spill���Բ���(Params)�����б��졣
    options.schedule_move_prob = 0.8;

    // [ʹ��˵��] �������ڱ���ʱ `min_farthest_use_threshold` �������Ա仯�������޷�Χ��
    options.threshold_range = { 100.0, 2000.0 };

    // --- �������ý��� ---
    // ========================================================================

    // --- �ļ����� ---
    const std::string nodes_file = std::string(DATA_DIR) + "/" + case_name + "_Nodes.csv";
    const std::string edges_file = std::string(DATA_DIR) + "/" + case_name + "_Edges.csv";
    const std::string schedule_file = std::string(DATA_DIR) + "/" + case_name + "_final_schedule_optimized.txt";

    Graph graph;
    CodeMap code_map;
    if (!FileIO::loadGraph(nodes_file, edges_file, graph, code_map)) {
        return 1;
    }
    std::cout << "Graph loaded." << std::endl;

    std::vector<int> p1_schedule;
    try {
        p1_schedule = FileIO::loadScheduleFromFile(schedule_file);
    }
    catch (const std::runtime_error& e) {
        std::cerr << e.what() << std::endl;
        return 1;
    }
    std::cout << "Schedule loaded, optimizing " << p1_schedule.size() << " nodes." << std::endl;

    // --- ִ���Ż� ---
    OptimizerNSGA optimizer(p1_schedule, graph, code_map, caps, params, options);
    std::vector<Individual> pareto_front = optimizer.run();

    // --- ��������뱣�� ---
    std::cout << "\n--- Optimization finished: Found " << pareto_front.size() << " Pareto optimal solutions ---" << std::endl;
    if (!pareto_front.empty()) {
        std::filesystem::path case_output_dir = std::string("output/Problem3_CPP/") + case_name;
        std::filesystem::create_directories(case_output_dir);

        std::vector<double> all_times, all_spills;
        for (const auto& sol : pareto_front) {
            all_times.push_back(sol.time);
            all_spills.push_back(sol.spill);
        }

        auto min_time_it = std::min_element(all_times.begin(), all_times.end());
        auto min_spill_it = std::min_element(all_spills.begin(), all_spills.end());
        int min_time_idx = std::distance(all_times.begin(), min_time_it);
        int min_spill_idx = std::distance(all_spills.begin(), min_spill_it);

        double min_t = *min_time_it;
        double range_t = *std::max_element(all_times.begin(), all_times.end()) - min_t;
        double min_s = *min_spill_it;
        double range_s = *std::max_element(all_spills.begin(), all_spills.end()) - min_s;

        int min_dist_idx = -1;
        double min_dist = std::numeric_limits<double>::infinity();
        for (int i = 0; i < pareto_front.size(); ++i) {
            double norm_t = (range_t > 0) ? (all_times[i] - min_t) / range_t : 0;
            double norm_s = (range_s > 0) ? (all_spills[i] - min_s) / range_s : 0;
            double dist = std::sqrt(norm_t * norm_t + norm_s * norm_s);
            if (dist < min_dist) {
                min_dist = dist;
                min_dist_idx = i;
            }
        }

        // === 保存帕累托前沿到CSV文件 ===
        std::cout << "\n--- Saving Pareto front to CSV for visualization ---" << std::endl;
        std::filesystem::path pareto_csv_path = case_output_dir / (case_name + "_pareto_solutions.csv");
        std::ofstream pareto_csv(pareto_csv_path);
        
        if (pareto_csv.is_open()) {
            // CSV 头部
            pareto_csv << "solution_id,time,spill,strategy_type,min_farthest_use_threshold,w_use,w_vol,w_vic" << std::endl;
            
            // 写入每个帕累托解的数据
            for (int i = 0; i < pareto_front.size(); ++i) {
                const auto& solution = pareto_front[i];
                
                pareto_csv << std::fixed << std::setprecision(2)
                          << (i + 1) << ","                                    // solution_id
                          << solution.time << ","                             // time
                          << solution.spill << ","                            // spill  
                          << solution.params.alloc_strategy << ","            // strategy_type
                          << solution.params.min_farthest_use_threshold << "," // threshold
                          << solution.params.w_use << ","                     // w_use
                          << solution.params.w_vol << ","                     // w_vol
                          << solution.params.w_vic                           // w_vic
                          << std::endl;
            }
            
            pareto_csv.close();
            std::cout << "  Pareto front saved to: " << pareto_csv_path << std::endl;
            std::cout << "  Total solutions: " << pareto_front.size() << std::endl;
        } else {
            std::cerr << "  ERROR: Could not open pareto CSV file for writing!" << std::endl;
        }

        std::cout << "\n--- Validating and saving all optimal solutions ---" << std::endl;
        for (int i = 0; i < pareto_front.size(); ++i) {
            const auto& current_solution_from_front = pareto_front[i];

            std::string folder_name = std::string("Solution_") + (i < 9 ? "0" : "") + std::to_string(i + 1);
            if (i == min_time_idx) folder_name += "_Best_Time";
            if (i == min_spill_idx) folder_name += "_Best_Spill";
            if (i == min_dist_idx) folder_name += "_Best_Tradeoff";

            std::cout << "\n---------------------------------------------------" << std::endl;
            std::cout << "--- Processing solution #" << (i + 1) << " (" << folder_name << ") ---" << std::endl;

            std::filesystem::path solution_dir = case_output_dir / folder_name;
            std::filesystem::create_directories(solution_dir);

            auto [solution_obj, final_graph] = CoreLogic::performAllocation(
                current_solution_from_front.schedule, graph, code_map, caps, current_solution_from_front.params);

            std::cout << "--- Validating solution for saving... ---" << std::endl;
            bool is_valid = CoreLogic::validateSolution(solution_obj, final_graph, code_map, caps);

            if (is_valid) {
                std::cout << "  Validation PASSED. Saving TXT files..." << std::endl;
                FileIO::writeScheduleToFile(solution_dir / (case_name + "_schedule.txt"), solution_obj.final_schedule);
                FileIO::writeMemoryToFile(solution_dir / (case_name + "_memory.txt"), solution_obj.memory_offsets);
                auto spill_log = FileIO::formatSpillLogForOutput(solution_obj.spill_map, final_graph);
                FileIO::writeSpillToFile(solution_dir / (case_name + "_spill.txt"), spill_log);
                std::cout << "  TXT results saved to: " << solution_dir << std::endl;
            }
            else {
                std::cerr << "!!! Validation FAILED: Solution #" << (i + 1) << " has issues. Files will not be saved. !!!" << std::endl;
            }
        }
    }

    std::cout << "\n--- All post-processing complete. ---" << std::endl;
    return 0;
}